//! 算术运算指令编码模块
//! 
//! 实现 8086 算术运算指令：ADD, SUB, CMP, MUL, DIV, INC, DEC, IMUL, IDIV

use crate::ast::Operand;
use crate::error::{AssemblerError, Result};
use super::registers::RegisterEncoding;

/// 算术运算指令编码器
pub struct ArithmeticEncoder;

impl ArithmeticEncoder {
    /// 编码算术运算指令
    pub fn encode(&self, instruction: &str, operands: &[Operand]) -> Result<Vec<u8>> {
        match instruction {
            "加" => self.encode_add(operands),
            "减" => self.encode_sub(operands),
            "比较" => self.encode_cmp(operands),
            "乘" => self.encode_mul(operands),
            "除" => self.encode_div(operands),
            "有符号乘" => self.encode_imul(operands),
            "有符号除" => self.encode_idiv(operands),
            "自增" => self.encode_inc(operands),
            "自减" => self.encode_dec(operands),
            _ => Err(AssemblerError::CodeGenError {
                message: format!("未知的算术指令: {}", instruction),
            }),
        }
    }

    /// 编码加指令
    fn encode_add(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 2 {
            return Err(AssemblerError::OperandError {
                message: "加 指令需要 2 个操作数".to_string(),
                line: 0,
            });
        }

        // ADD reg, reg: 0000000w | mod reg r/m
        // ADD reg, imm: 100000sw | mod 000 r/m | imm
        match (&operands[0], &operands[1]) {
            (Operand::Register(dst), Operand::Register(src)) => {
                let dst_code = RegisterEncoding::get_16bit_register_code(dst)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", dst),
                        line: 0,
                    })?;
                let src_code = RegisterEncoding::get_16bit_register_code(src)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", src),
                        line: 0,
                    })?;
                let opcode = 0x01; // ADD reg, reg (16位)
                let modrm = (0b11 << 6) | (dst_code << 3) | src_code;
                Ok(vec![opcode, modrm])
            }
            (Operand::Register(reg), Operand::Immediate(imm)) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                if *imm >= -128 && *imm <= 127 {
                    // 8位立即数
                    let opcode = 0x83; // ADD reg, imm8
                    let modrm = (0b11 << 6) | (0b000 << 3) | reg_code;
                    Ok(vec![opcode, modrm, *imm as u8])
                } else {
                    // 16位立即数
                    let opcode = 0x81; // ADD reg, imm16
                    let modrm = (0b11 << 6) | (0b000 << 3) | reg_code;
                    let imm_u16 = *imm as u16;
                    Ok(vec![opcode, modrm, imm_u16 as u8, (imm_u16 >> 8) as u8])
                }
            }
            _ => Err(AssemblerError::OperandError {
                message: "加暂不支持此操作数组合".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码减指令
    fn encode_sub(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 2 {
            return Err(AssemblerError::OperandError {
                message: "减 指令需要 2 个操作数".to_string(),
                line: 0,
            });
        }

        match (&operands[0], &operands[1]) {
            (Operand::Register(dst), Operand::Register(src)) => {
                let dst_code = RegisterEncoding::get_16bit_register_code(dst)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", dst),
                        line: 0,
                    })?;
                let src_code = RegisterEncoding::get_16bit_register_code(src)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", src),
                        line: 0,
                    })?;
                let opcode = 0x29; // SUB reg, reg (16位)
                let modrm = (0b11 << 6) | (dst_code << 3) | src_code;
                Ok(vec![opcode, modrm])
            }
            (Operand::Register(reg), Operand::Immediate(imm)) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                if *imm >= -128 && *imm <= 127 {
                    let opcode = 0x83; // SUB reg, imm8
                    let modrm = (0b11 << 6) | (0b101 << 3) | reg_code;
                    Ok(vec![opcode, modrm, *imm as u8])
                } else {
                    let opcode = 0x81; // SUB reg, imm16
                    let modrm = (0b11 << 6) | (0b101 << 3) | reg_code;
                    let imm_u16 = *imm as u16;
                    Ok(vec![opcode, modrm, imm_u16 as u8, (imm_u16 >> 8) as u8])
                }
            }
            _ => Err(AssemblerError::OperandError {
                message: "减暂不支持此操作数组合".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码比较指令
    fn encode_cmp(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 2 {
            return Err(AssemblerError::OperandError {
                message: "比较 指令需要 2 个操作数".to_string(),
                line: 0,
            });
        }

        // CMP reg, reg: 0011100w | mod reg r/m
        // CMP reg, imm: 100000sw | mod 111 r/m | imm
        match (&operands[0], &operands[1]) {
            (Operand::Register(dst), Operand::Register(src)) => {
                let dst_code = RegisterEncoding::get_16bit_register_code(dst)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", dst),
                        line: 0,
                    })?;
                let src_code = RegisterEncoding::get_16bit_register_code(src)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", src),
                        line: 0,
                    })?;
                let opcode = 0x39; // CMP reg, reg (16位)
                let modrm = (0b11 << 6) | (dst_code << 3) | src_code;
                Ok(vec![opcode, modrm])
            }
            (Operand::Register(reg), Operand::Immediate(imm)) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                if *imm >= -128 && *imm <= 127 {
                    // 8位立即数
                    let opcode = 0x83; // CMP reg, imm8
                    // reg字段=111 (0b111) 表示CMP
                    let modrm = (0b11 << 6) | (0b111 << 3) | reg_code;
                    Ok(vec![opcode, modrm, *imm as u8])
                } else {
                    // 16位立即数
                    let opcode = 0x81; // CMP reg, imm16
                    // reg字段=111 (0b111) 表示CMP
                    let modrm = (0b11 << 6) | (0b111 << 3) | reg_code;
                    let imm_u16 = *imm as u16;
                    Ok(vec![opcode, modrm, imm_u16 as u8, (imm_u16 >> 8) as u8])
                }
            }
            _ => Err(AssemblerError::OperandError {
                message: "比较暂不支持此操作数组合".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码乘法指令（MUL）
    fn encode_mul(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: "乘 指令需要 1 个操作数".to_string(),
                line: 0,
            });
        }

        match &operands[0] {
            Operand::Register(reg) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                // MUL reg16：F7 /4 (reg字段=100表示MUL)
                let opcode = 0xF7;
                let modrm = (0b11 << 6) | (0b100 << 3) | reg_code;
                Ok(vec![opcode, modrm])
            }
            _ => Err(AssemblerError::OperandError {
                message: "乘 指令只支持寄存器操作数".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码除法指令（DIV）
    fn encode_div(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: "除 指令需要 1 个操作数".to_string(),
                line: 0,
            });
        }

        match &operands[0] {
            Operand::Register(reg) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                // DIV reg16：F7 /6 (reg字段=110表示DIV)
                let opcode = 0xF7;
                let modrm = (0b11 << 6) | (0b110 << 3) | reg_code;
                Ok(vec![opcode, modrm])
            }
            _ => Err(AssemblerError::OperandError {
                message: "除 指令只支持寄存器操作数".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码有符号乘法指令（IMUL）
    /// 
    /// IMUL reg16：F7 /5 (reg字段=101表示IMUL)
    /// 操作：AX = AX * reg (有符号乘法)
    /// 结果：DX:AX = 32位有符号结果
    fn encode_imul(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: "有符号乘 指令需要 1 个操作数".to_string(),
                line: 0,
            });
        }

        match &operands[0] {
            Operand::Register(reg) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                // IMUL reg16：F7 /5 (reg字段=101表示IMUL)
                let opcode = 0xF7;
                let modrm = (0b11 << 6) | (0b101 << 3) | reg_code;
                Ok(vec![opcode, modrm])
            }
            _ => Err(AssemblerError::OperandError {
                message: "有符号乘 指令只支持寄存器操作数".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码有符号除法指令（IDIV）
    /// 
    /// IDIV reg16：F7 /7 (reg字段=111表示IDIV)
    /// 操作：DX:AX / reg (有符号除法)
    /// 结果：AX = 商，DX = 余数
    fn encode_idiv(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: "有符号除 指令需要 1 个操作数".to_string(),
                line: 0,
            });
        }

        match &operands[0] {
            Operand::Register(reg) => {
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                // IDIV reg16：F7 /7 (reg字段=111表示IDIV)
                let opcode = 0xF7;
                let modrm = (0b11 << 6) | (0b111 << 3) | reg_code;
                Ok(vec![opcode, modrm])
            }
            _ => Err(AssemblerError::OperandError {
                message: "有符号除 指令只支持寄存器操作数".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码自增指令
    fn encode_inc(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: "自增 指令需要 1 个操作数".to_string(),
                line: 0,
            });
        }

        if let Operand::Register(reg) = &operands[0] {
            let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                .ok_or_else(|| AssemblerError::OperandError {
                    message: format!("无效的16位寄存器: {}", reg),
                    line: 0,
                })?;
            let opcode = 0x40 + reg_code; // INC reg (16位)
            Ok(vec![opcode])
        } else {
            Err(AssemblerError::OperandError {
                message: "自增只支持寄存器操作数".to_string(),
                line: 0,
            })
        }
    }

    /// 编码自减指令
    fn encode_dec(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: "自减 指令需要 1 个操作数".to_string(),
                line: 0,
            });
        }

        if let Operand::Register(reg) = &operands[0] {
            let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                .ok_or_else(|| AssemblerError::OperandError {
                    message: format!("无效的16位寄存器: {}", reg),
                    line: 0,
                })?;
            let opcode = 0x48 + reg_code; // DEC reg (16位)
            Ok(vec![opcode])
        } else {
            Err(AssemblerError::OperandError {
                message: "自减只支持寄存器操作数".to_string(),
                line: 0,
            })
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encode_add_reg_reg() {
        let encoder = ArithmeticEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("加", &operands).unwrap();
        // ADD AX, BX = 01 C3 (dst=AX=0, src=BX=3, modrm=0b11000011=0xC3)
        assert_eq!(result, vec![0x01, 0xC3]);
    }

    #[test]
    fn test_encode_add_reg_imm8() {
        let encoder = ArithmeticEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(10),
        ];
        let result = encoder.encode("加", &operands).unwrap();
        // ADD AX, 10 = 83 C0 0A
        assert_eq!(result, vec![0x83, 0xC0, 0x0A]);
    }

    #[test]
    fn test_encode_add_reg_imm16() {
        let encoder = ArithmeticEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Immediate(1000),
        ];
        let result = encoder.encode("加", &operands).unwrap();
        // ADD AX, 1000 = 81 C0 E8 03
        assert_eq!(result, vec![0x81, 0xC0, 0xE8, 0x03]);
    }

    #[test]
    fn test_encode_sub_reg_reg() {
        let encoder = ArithmeticEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("减", &operands).unwrap();
        // SUB AX, BX = 29 C3 (dst=AX=0, src=BX=3, modrm=0b11000011=0xC3)
        assert_eq!(result, vec![0x29, 0xC3]);
    }

    #[test]
    fn test_encode_inc() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("累加器".to_string())];
        let result = encoder.encode("自增", &operands).unwrap();
        // INC AX = 40
        assert_eq!(result, vec![0x40]);
    }

    #[test]
    fn test_encode_dec() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("基址".to_string())];
        let result = encoder.encode("自减", &operands).unwrap();
        // DEC BX = 4B
        assert_eq!(result, vec![0x4B]);
    }

    #[test]
    fn test_encode_cmp_reg_reg() {
        let encoder = ArithmeticEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("比较", &operands).unwrap();
        // CMP AX, BX = 39 C3 (dst=AX=0, src=BX=3, modrm=0b11000011=0xC3)
        assert_eq!(result, vec![0x39, 0xC3]);
    }

    #[test]
    fn test_encode_mul() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("累加器".to_string())];
        let result = encoder.encode("乘", &operands).unwrap();
        // MUL AX = F7 E0
        assert_eq!(result, vec![0xF7, 0xE0]);
    }

    #[test]
    fn test_encode_div() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("基址".to_string())];
        let result = encoder.encode("除", &operands).unwrap();
        // DIV BX = F7 F3
        assert_eq!(result, vec![0xF7, 0xF3]);
    }

    #[test]
    fn test_encode_invalid_instruction() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("累加器".to_string())];
        let result = encoder.encode("无效指令", &operands);
        assert!(result.is_err());
    }

    #[test]
    fn test_encode_add_wrong_operand_count() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("累加器".to_string())];
        let result = encoder.encode("加", &operands);
        assert!(result.is_err());
    }

    #[test]
    fn test_encode_imul() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("累加器".to_string())];
        let result = encoder.encode("有符号乘", &operands).unwrap();
        // IMUL AX = F7 E8 (reg字段=101=0b101, AX=000, modrm=0b11101000=0xE8)
        assert_eq!(result, vec![0xF7, 0xE8]);
    }

    #[test]
    fn test_encode_imul_bx() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("基址".to_string())];
        let result = encoder.encode("有符号乘", &operands).unwrap();
        // IMUL BX = F7 EB (reg字段=101=0b101, BX=011, modrm=0b11101011=0xEB)
        assert_eq!(result, vec![0xF7, 0xEB]);
    }

    #[test]
    fn test_encode_idiv() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("基址".to_string())];
        let result = encoder.encode("有符号除", &operands).unwrap();
        // IDIV BX = F7 FB (reg字段=111=0b111, BX=011, modrm=0b11111011=0xFB)
        assert_eq!(result, vec![0xF7, 0xFB]);
    }

    #[test]
    fn test_encode_idiv_ax() {
        let encoder = ArithmeticEncoder;
        let operands = vec![Operand::Register("累加器".to_string())];
        let result = encoder.encode("有符号除", &operands).unwrap();
        // IDIV AX = F7 F8 (reg字段=111=0b111, AX=000, modrm=0b11111000=0xF8)
        assert_eq!(result, vec![0xF7, 0xF8]);
    }

    #[test]
    fn test_encode_imul_wrong_operand_count() {
        let encoder = ArithmeticEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("有符号乘", &operands);
        assert!(result.is_err());
    }

    #[test]
    fn test_encode_idiv_wrong_operand_count() {
        let encoder = ArithmeticEncoder;
        let operands = vec![
            Operand::Register("累加器".to_string()),
            Operand::Register("基址".to_string()),
        ];
        let result = encoder.encode("有符号除", &operands);
        assert!(result.is_err());
    }
}

